package org.tigr.htc.common.sge;

import com.sun.grid.jgdi.JGDI;
import com.sun.grid.jgdi.JGDIException;
import com.sun.grid.jgdi.JGDIFactory;
import com.sun.grid.jgdi.configuration.ClusterQueue;
import com.sun.grid.jgdi.configuration.ClusterQueueImpl;
import com.sun.grid.jgdi.util.shell.editor.EditorUtil;
import org.tigr.htc.common.sge.JGDILocator;

public class SgeQueueConfigurator {

    private static final String ARGS = "$host $job_owner $job_id $job_name $queue";
    private String htcRoot;
    private JGDI jgdi;
    public static final String SHELL_BEHAVIOR = "unix_behavior";
    private static final String DEFAULT = EditorUtil.VALUE_DEFAULT;
    private static final String NONE = EditorUtil.VALUE_NONE;
    private static final String UNDEFINED = EditorUtil.VALUE_UNDEFINED;
    private static final String INFINITY = EditorUtil.VALUE_INFINITY;

    /**
     * @param htcRoot
     */
    public SgeQueueConfigurator(JGDI jgdi, String htcRoot) {
        this.htcRoot = htcRoot;
        this.jgdi = jgdi;
    }

    /**
     * @param queueName
     * @param shell
     * @throws JGDIException
     */
    public void configureQueue(String queueName, String shell) throws JGDIException {
        try {
            ClusterQueue queue = jgdi.getClusterQueue(queueName);
            queue.putEpilog(DEFAULT, htcRoot + "/bin/epilog " + ARGS);
            queue.putProlog(DEFAULT, htcRoot + "/bin/prolog " + ARGS);
            queue.putShellStartMode(DEFAULT, SHELL_BEHAVIOR);
            queue.putShell(DEFAULT, shell);

            // Update the queue with the new queue configuration
            jgdi.updateClusterQueue(queue);
        } catch (JGDIException je) {
            je.printStackTrace();
            throw je;
        }
    }

    public void createQueue(String queueName, String shell) throws Exception {
        try {
            ClusterQueue queue = new ClusterQueueImpl(queueName);
            queue.putJobSlots(DEFAULT, 0);
            queue.putSeqNo(DEFAULT, 0);
            queue.putNsuspend(DEFAULT, 1);
            queue.putTmpdir(DEFAULT, "/tmp");
            queue.putShell(DEFAULT, shell);
            queue.putCalendar(DEFAULT, NONE);
            queue.putPriority(DEFAULT, "0");
            queue.putProcessors(DEFAULT, UNDEFINED);
            queue.putProlog(DEFAULT, htcRoot + "/bin/prolog " + ARGS);
            queue.putEpilog(DEFAULT, htcRoot + "/bin/epilog " + ARGS);
            queue.putShellStartMode(DEFAULT, SHELL_BEHAVIOR);
            queue.putStarterMethod(DEFAULT, NONE);
            queue.putSuspendMethod(DEFAULT, NONE);
            queue.putResumeMethod(DEFAULT, NONE);
            queue.putTerminateMethod(DEFAULT, NONE);
            queue.putInitialState(DEFAULT, "default");
            queue.putRerun(DEFAULT, false);
            queue.putSFsize(DEFAULT, INFINITY);
            queue.putHFsize(DEFAULT, INFINITY);
            queue.putSData(DEFAULT, INFINITY);
            queue.putHData(DEFAULT, INFINITY);
            queue.putSStack(DEFAULT, INFINITY);
            queue.putHStack(DEFAULT, INFINITY);
            queue.putSCore(DEFAULT, INFINITY);
            queue.putHCore(DEFAULT, INFINITY);
            queue.putSRss(DEFAULT, INFINITY);
            queue.putHRss(DEFAULT, INFINITY);
            queue.putSVmem(DEFAULT, INFINITY);
            queue.putHVmem(DEFAULT, INFINITY);
            queue.putSRt(DEFAULT, INFINITY);
            queue.putHRt(DEFAULT, INFINITY);
            queue.putSCpu(DEFAULT, INFINITY);
            queue.putHCpu(DEFAULT, INFINITY);
            queue.putSuspendInterval(DEFAULT, "00:05:00");
            queue.putMinCpuInterval(DEFAULT, "00:05:00");
            queue.putNotify(DEFAULT, "00:00:60");
            // TODO, there is a method in EditorUtil to go from "BATCH", the String form, to the Integer...
            queue.putQtype(DEFAULT, 3);

            queue.addEmptyCkpt(DEFAULT);
            queue.addEmptyPe(DEFAULT);
            queue.addEmptyOwner(DEFAULT);
            queue.addEmptyAcl(DEFAULT);
            queue.addEmptyXacl(DEFAULT);
            queue.addEmptyProjects(DEFAULT);
            queue.addEmptyXprojects(DEFAULT);
            queue.addEmptyConsumableConfig(DEFAULT);
            queue.addEmptyLoadThresholds(DEFAULT);
            queue.addEmptySuspendThresholds(DEFAULT);
            queue.addEmptySubordinate(DEFAULT);
            queue.addHostlist("@allhosts");

            jgdi.addClusterQueue(queue);
        } catch (Exception je) {
            je.printStackTrace();
            throw je;
        }    
    }

    /**
     * @param queueName
     * @return true or false
     * @throws JGDIException
     */
    public boolean queueExists(String queueName) throws JGDIException {
        boolean exists = false;
        try {
            ClusterQueue queue = jgdi.getClusterQueue(queueName);
            if (queue != null) {
                exists = true;
            }
        } catch (JGDIException je) {
            je.printStackTrace();
            throw je;
        }
        return exists;
    }

    public static void main(String[] args) {
        JGDI myjgdi = null;

        if (args.length != 3) {
            System.err.println("3 arguments required: <HTC root> <HTC queue name> <queue shell>");
            System.exit(1);
        }
        String url = args[0];
        String workflowRoot = url.replaceFirst("bootstrap://", "").replaceFirst("@.*", "");
        System.out.println(workflowRoot);
        String queueName = args[1];
        String shell = args[2];

        try {
            myjgdi = JGDILocator.getJGDI();
            SgeQueueConfigurator configurator = new SgeQueueConfigurator(myjgdi, workflowRoot);

            if (configurator.queueExists(queueName)) {
                System.out.println("Queue " + queueName + " already exists. Configuring it.");
                configurator.configureQueue(queueName, shell);
            } else {
                System.out.println("Queue " + queueName + " doesn't exist. Creating it.");
                configurator.createQueue(queueName, shell);
            }
        } catch (Exception e) {
            e.printStackTrace();
            System.exit(1);
        } finally {
            if (myjgdi != null) {
                try {
                    myjgdi.close();
                } catch (JGDIException ex) {
                    ex.printStackTrace();
                }
            }
        }
    }
}
